feat(spanner): support for Multiplexed Session Partitioned Ops (#2252)

* feat(spanner): support for Multiplexed Session Partitioned Ops

* feat(spanner): add support for snapshot isolation (#2245)

This PR contains code changes to add support for option IsolationLevel at the client level and at the transaction level.
supported methods are(RW and Blind Write):

```
- writeAtLeastOnce
- runTransactionAsync
- runTransaction
- getTransaction
- async getTransaction(from transaction runner class)
```

* refactor

* refactor

Author: alkatrivedi

.kokoro/presubmit/node14/system-test-multiplexed-session.cfg

@@ -1964,7 +1964,7 @@ describe('Database', () => {
       },
     };
 
-    let fakePool: FakeSessionPool;
+    let fakeSessionFactory: FakeSessionFactory;
     let fakeSession: FakeSession;
     let fakePartitionedDml = new FakeTransaction(
       {} as google.spanner.v1.TransactionOptions.PartitionedDml
@@ -1974,14 +1974,17 @@ describe('Database', () => {
     let beginStub;
 
     beforeEach(() => {
-      fakePool = database.pool_;
+      fakeSessionFactory = database.sessionFactory_;
       fakeSession = new FakeSession();
       fakePartitionedDml = new FakeTransaction(
         {} as google.spanner.v1.TransactionOptions.PartitionedDml
       );
 
       getSessionStub = (
-        sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
+        sandbox.stub(
+          fakeSessionFactory,
+          'getSessionForPartitionedOps'
+        ) as sinon.SinonStub
       ).callsFake(callback => {
         callback(null, fakeSession);
       });
@@ -2077,7 +2080,7 @@ describe('Database', () => {
       beginStub.callsFake(callback => callback(fakeError));
 
       const releaseStub = (
-        sandbox.stub(fakePool, 'release') as sinon.SinonStub
+        sandbox.stub(fakeSessionFactory, 'release') as sinon.SinonStub
       ).withArgs(fakeSession);
 
       database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
@@ -2124,7 +2127,7 @@ describe('Database', () => {
 
     it('session released on transaction end', done => {
       const releaseStub = (
-        sandbox.stub(fakePool, 'release') as sinon.SinonStub
+        sandbox.stub(fakeSessionFactory, 'release') as sinon.SinonStub
       ).withArgs(fakeSession);
 
       database.runPartitionedUpdate(QUERY, async () => {

.kokoro/trampoline_v2.sh

@@ -2945,7 +2945,7 @@ class Database extends common.GrpcServiceObject {
       ...this._traceConfig,
     };
     return startTrace('Database.runPartitionedUpdate', traceConfig, span => {
-      this.pool_.getSession((err, session) => {
+      this.sessionFactory_.getSessionForPartitionedOps((err, session) => {
         if (err) {
           setSpanError(span, err);
           span.end();
@@ -2976,21 +2976,21 @@ class Database extends common.GrpcServiceObject {
     }
     transaction.begin(err => {
       if (err) {
-        this.pool_.release(session!);
+        this.sessionFactory_.release(session!);
         callback!(err, 0);
         return;
       }
 
       transaction.runUpdate(query, (err, updateCount) => {
         if (err) {
           if (err.code !== grpc.status.ABORTED) {
-            this.pool_.release(session!);
+            this.sessionFactory_.release(session!);
             callback!(err, 0);
             return;
           }
           this._runPartitionedUpdate(session, query, callback);
         } else {
-          this.pool_.release(session!);
+          this.sessionFactory_.release(session!);
           callback!(null, updateCount);
           return;
         }

observability-test/database.ts

@@ -56,6 +56,14 @@ export interface SessionFactoryInterface {
    */
   getSession(callback: GetSessionCallback): void;
 
+  /**
+   * When called returns a session for paritioned dml.
+   *
+   * @name SessionFactoryInterface#getSessionForPartitionedOps
+   * @param {GetSessionCallback} callback The callback function.
+   */
+  getSessionForPartitionedOps(callback: GetSessionCallback): void;
+
   /**
    * When called returns the pool object.
    *
@@ -97,6 +105,7 @@ export class SessionFactory
   multiplexedSession_: MultiplexedSessionInterface;
   pool_: SessionPoolInterface;
   isMultiplexed: boolean;
+  isMultiplexedPartitionOps: boolean;
   constructor(
     database: Database,
     name: String,
@@ -117,6 +126,11 @@ export class SessionFactory
     process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true'
       ? (this.isMultiplexed = true)
       : (this.isMultiplexed = false);
+    // set the isMultiplexedPartitionedOps property to true if multiplexed session is enabled for paritioned ops, otherwise set the property to false
+    this.isMultiplexedPartitionOps =
+      process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS === 'true' &&
+      process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS ===
+        'true';
     // Multiplexed sessions should only be created if its enabled.
     if (this.isMultiplexed) {
       this.multiplexedSession_.on('error', this.emit.bind(database, 'error'));
@@ -143,6 +157,23 @@ export class SessionFactory
     );
   }
 
+  /**
+   * Retrieves a session for partitioned operations, selecting the appropriate session type
+   * based on whether multiplexed sessions are enabled.
+   *
+   * If multiplexed sessions are enabled for partitioned ops this methods delegates the request to `getSession()`, which returns
+   * either a multiplexed session or a regular session based on the configuration.
+   *
+   * If the multiplexed sessions are disabled, a session is retrieved from the regular session pool.
+   *
+   * @param {GetSessionCallback} callback The callback function.
+   */
+  getSessionForPartitionedOps(callback: GetSessionCallback): void {
+    this.isMultiplexedPartitionOps
+      ? this.getSession(callback)
+      : this.pool_.getSession(callback);
+  }
+
   /**
    * Returns the regular session pool object.
    *
@@ -154,9 +185,11 @@ export class SessionFactory
   }
 
   /**
-   * Releases a session back to the session pool.
+   * Releases a regular session back to the session pool.
+   *
+   * This methods does not release a multiplexed session.
    *
-   * This method returns a session to the pool after it is no longer needed.
+   * It returns a session to the pool after it is no longer needed.
    * It is a no-op for multiplexed sessions.
    *
    * @param {Session} session - The session to be released. This should be an instance of `Session` that was
@@ -165,7 +198,7 @@ export class SessionFactory
    * @throws {Error} If the session is invalid or cannot be released.
    */
   release(session: Session): void {
-    if (!this.isMultiplexed) {
+    if (!session.metadata?.multiplexed) {
       this.pool_.release(session);
     }
   }